import sqlparse
from sqlparse.sql import TokenList
import numpy as np
import seaborn as sns
import  pandas as pd
from tqdm import  tqdm
import matplotlib.pyplot as plt

# 设置matplotlib字体
# plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体
plt.rcParams['font.sans-serif'] = ['Heiti TC']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

depth_list = []
width_list = []
def get_tree_depth_width(node, current_depth=0):
    # 基本情况：如果节点不是一个 TokenList，那么它是一个叶节点
    if not isinstance(node, TokenList):
        return current_depth

    # 递归地计算所有深度的宽度
    max_depth = current_depth
    len = 0
    for item in node:
        if str(item) != ' ':
           len += 1
    width_list[current_depth] += len

    # 递归地计算所有子节点的深度
    for item in node:
        depth = get_tree_depth_width(item, current_depth + 1)
        max_depth = max(max_depth, depth)

    return max_depth

def get_sqls_complexity(sql_path):
    global  depth_list, width_list
    depth_list = [0] * 50
    width_list = [0] * 50

    df = pd.read_csv(sql_path)
    sql_complexity = df[['SQL语句']]
    sql_complexity = sql_complexity[sql_complexity['SQL语句'].apply(lambda x: x.strip()[:6].upper() == 'SELECT')]
    depth_value_list = []
    width_value_list = []
    for sql in tqdm(sql_complexity['SQL语句'].values):
        width_list = [0] * 50
        sql = sql.strip()
        if (len(sql) != 0):
            # 使用 sqlparse 解析 SQL 语句
            parsed = sqlparse.parse(sql)[0]
            # 计算树的深度和宽度
            depth = get_tree_depth_width(parsed)
            depth_value_list.append(depth)
            width_value_list.append(max(width_list))
    sql_complexity['AST_depth'] = depth_value_list
    sql_complexity['AST_width'] = width_value_list
    return sql_complexity

def visualization(sql_complexity, figure_save_path, column):
    plt.figure(figsize=(15, 6))
    # 绘制直方图并获取柱子对象
    if column.endswith("depth"):
        bin_edges = list(np.arange(1, 21)) + [21, 26, 31, 36, 41, 46, 50]  # 自定义分桶边界
    if column.endswith("width"):
        # bin_edges = list(np.arange(0, 100, 5)) + [200, 300, 400, 500, 3000]  # 自定义分桶边界
        bin_edges = np.arange(1, 100, 5)

    ax = sns.histplot(sql_complexity[column], bins=bin_edges, palette='bright')

    # 在每个柱子上显示数量
    for p in ax.patches:
        ax.text(p.get_x() + p.get_width() / 2., p.get_height(), '%d' % int(p.get_height()),
                fontsize=12, color='red', ha='center', va='bottom')
    plt.tight_layout()
    plt.savefig(figure_save_path, format='png', dpi=300)
    # 显示图表
    plt.show()


def get_complexity_distribution(sql_path, depth_distribution_save_path,depth_figure_save_path,width_distribution_save_path,width_figure_save_path):
    sql_complexity = get_sqls_complexity(sql_path)
    complexity_dict = {"AST_depth":"深度", "AST_width":"宽度" }
    for key, value in complexity_dict.items():
        complexity_distribution = pd.DataFrame(sql_complexity[key].value_counts())
        complexity_distribution[value] = complexity_distribution.index
        complexity_distribution.sort_values(by=value, axis=0, inplace=True)
        complexity_distribution.rename(columns = {key: "SQL数量"}, inplace=True)
        complexity_distribution['SQL数量占比'] = complexity_distribution['SQL数量']/sql_complexity.shape[0] * 100
        complexity_distribution['SQL数量占比'] = complexity_distribution['SQL数量占比'].apply(lambda x: f'{x:.2f}%')
        if key.endswith("depth"):
            complexity_distribution[[value, 'SQL数量','SQL数量占比']].to_csv(depth_distribution_save_path, index=False)
        else:
            complexity_distribution[[value, 'SQL数量','SQL数量占比']].to_csv(width_distribution_save_path, index=False)

    visualization(sql_complexity, depth_figure_save_path, "AST_depth")
    visualization(sql_complexity, width_figure_save_path, "AST_width")


if __name__ == "__main__":
    get_complexity_distribution("../data/电力/单表问答对.csv",
                        "../data/电力/electricity_sql_complexity_depth.csv",
                        "../data/电力/electricity_sql_complexity_depth.png",
                        "../data/电力/electricity_sql_complexity_width.csv",
                        "../data/电力/electricity_sql_complexity_width.png")

    get_complexity_distribution("../data/浪潮/问答对.csv",
                        "../data/浪潮/inspur_sql_complexity_depth.csv",
                        "../data/浪潮/inspur_sql_complexity_depth.png",
                        "../data/浪潮/inspur_sql_complexity_width.csv",
                        "../data/浪潮/inspur_sql_complexity_width.png")
